"""
@Time ： 2025/11/26 16:25
@Auth ： seven
@File ：utils.py
@IDE ：PyCharm
"""
import os
import json
import datetime
import torch
from collections import OrderedDict


def save_experiment_config(opt, save_dir):
    """保存精简的实验配置参数"""
    # 创建保存目录
    os.makedirs(save_dir, exist_ok=True)

    # 只记录最重要的参数
    config = {
        # 基本信息
        'experiment_name': opt.name,
        'timestamp': datetime.datetime.now().strftime('%Y-%m-%d_%H-%M-%S'),
        'device': 'GPU: ' + torch.cuda.get_device_name(0) if torch.cuda.is_available() else 'CPU',

        # 训练配置
        'training_config': {
            'batch_size': opt.batchsize,
            'epochs': opt.niter,
            'learning_rate': opt.lr,
            'beta1': opt.beta1,
            'dataset': opt.dataset,
            'normal_class': opt.nc,
            'use_semi': opt.use_semi,
            'use_lr_scheduler': getattr(opt, 'use_lr_scheduler', False)
        },

        # 模型结构
        'model_config': {
            'latent_dim': opt.nz,
            'generator_filters': opt.ngf,
            'discriminator_filters': opt.ndf,
            'extra_layers': opt.extralayers
        },

        # 损失权重（核心参数）
        'loss_weights': {
            'w_adv': opt.w_adv,
            'w_con': opt.w_con,
            'w_enc': opt.w_enc,
            'w_cls': getattr(opt, 'w_cls', 0),
            'w_contrast': getattr(opt, 'w_contrast', 0),
            'w_push': getattr(opt, 'w_push', 0),
            # 启用的损失标志
            'use_contrast': getattr(opt, 'use_contrast', False),
            'use_push': getattr(opt, 'use_push', False),
            'use_latent_reg': getattr(opt, 'use_latent_reg', False)
        }
    }

    # 生成文件名
    filename = f"config_{opt.name}_{config['timestamp'].replace(':', '-')}.json"
    filepath = os.path.join(save_dir, filename)

    # 保存为JSON文件
    with open(filepath, 'w', encoding='utf-8') as f:
        json.dump(config, f, indent=4, ensure_ascii=False)

    # 同时生成一个简洁的txt版本（便于快速查看）
    txt_filename = f"config_summary_{opt.name}.txt"
    txt_filepath = os.path.join(save_dir, txt_filename)
    with open(txt_filepath, 'w', encoding='utf-8') as f:
        f.write(f"Experiment: {opt.name}\n")
        f.write(f"Date: {config['timestamp']}\n")
        f.write("=" * 50 + "\n\n")

        f.write("Training Config:\n")
        for k, v in config['training_config'].items():
            f.write(f"  {k}: {v}\n")

        f.write("\nModel Config:\n")
        for k, v in config['model_config'].items():
            f.write(f"  {k}: {v}\n")

        f.write("\nLoss Weights:\n")
        for k, v in config['loss_weights'].items():
            f.write(f"  {k}: {v}\n")

    print(f"实验配置已保存到: {filepath}")
    return filepath